from tensorflow.keras.layers import Layer
from tensorflow.keras import backend as K
import tensorflow as tf


class BinaryRandomChoice(Layer):
    """随机二选一
        """
    def __init__(self,**kwargs):
        super(BinaryRandomChoice,self).__init__(**kwargs)
        self.supports_masking = True

    def compute_mask(self,inputs,mask=None):
        if mask is not None:
            return mask[1]

    def call(self,inputs):
        source,target = inputs
        mask = K.random_binomial(shape=[1],p=0.5)
        output = mask * source + (1 - mask) * target
        return K.in_train_phase(output,target)

    def compute_output_shape(self, input_shape):
        return input_shape[1]

class BahdanauAttentionLayer(Layer):
    """
        This class implements Bahdanau attention (https://arxiv.org/pdf/1409.0473.pdf).
        There are three sets of weights introduced W_a, U_a, and V_a
         """

    def __init__(self,units):
        super(BahdanauAttentionLayer, self).__init__()
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)



    def call(self,query,values):
        # 隐藏层的形状 == （批大小，隐藏层大小）
        # hidden_with_time_axis 的形状 == （批大小，1，隐藏层大小）
        # 这样做是为了执行加法以计算分数
        hidden_with_time_axis = tf.expand_dims(query,1)

        # 分数的形状 == （批大小，最大长度，1）
        # 我们在最后一个轴上得到 1， 因为我们把分数应用于 self.V
        # 在应用 self.V 之前，张量的形状是（批大小，最大长度，单位）
        score = self.V(tf.nn.tanh(
            self.W1(values) + self.W2(hidden_with_time_axis)
        ))
        # 注意力权重 （attention_weights） 的形状 == （批大小，最大长度，1）
        attention_weights = tf.nn.softmax(score,axis=1)

        # 上下文向量 （context_vector） 求和之后的形状 == （批大小，隐藏层大小
        context_vector = attention_weights * values
        context_vector = tf.reduce_sum(context_vector,axis=1)

        return context_vector,attention_weights


if __name__ == '__main__':
    pass
